from py2neo import Node, Relationship, Graph, NodeMatcher, RelationshipMatcher
import csv


print("输入完成")

# Neo4j数据库连接信息
url = "bolt://localhost:7687"
username = "neo4j"
password = "00000000"
graph = Graph(url, auth=(username, password))


def find_or_create_node(label, name):
    # 尝试找到节点
    query = f"MATCH (n:{label}" + " {name: '" + name + "'}) RETURN COUNT(n) AS count;"
    result = graph.run(query).data()[0]['count']
    new_node = Node(label, name=name)
    # print(f"node: {name}, count: {result}\n") # debug
    # 如果没找到，则创建
    if result == 0:
        graph.create(new_node)
    # return new_node
# 主函数
def main():
    # node1 = Node('person', name='张三')
    # node2 = Node('address', name='北京市海淀区')
    # node3 = Node('person', name='李四')
    #
    # # 结点实例化
    # graph.create(node1)
    # graph.create(node2)
    # graph.create(node3)
    #
    # # 创建关系
    # address = Relationship(node1, '地址', node2)
    # classmates = Relationship(node1, '同学', node3)
    # address1 = Relationship(node3, '地址', node2)
    #
    # # 关系实例化
    # graph.create(address)
    # graph.create(address1)
    # graph.create(classmates)

    # row_list = [['id', 'name', 'age', 'sex', 'address'],
    #             [0, 'John', '23', 'man', 'Beijing'],
    #             [1, 'Bob', '23', 'man', 'Shanghai'],
    #             [2, 'Ben', '23', 'man', 'Shanghai'],
    #             [3, 'Mike', '23', 'man', 'Beijing'],
    #             [4, 'John', '23', 'man', 'Beijing'],
    #             [5, 'Robert J', '23', 'man', 'Beijing'],
    #             [6, 'Zhangsan', '23', 'man', 'ShanDong'],
    #             [7, 'Lisisi', '23', 'man', 'ShanDong'],
    #             [8, 'Liuyiyi', '23', 'man', 'Beijing'],
    #             [9, 'Heyi', '23', 'man', 'JiangSu'], ]

    # with open('data.csv', 'w', newline='') as f:
    #     writer = csv.writer(f)
    #     writer.writerows(row_list)
    #
    # print("输入完成")
    # 先清空库
    graph.delete_all()
    print("Data deleted successfully.")

    # region 创建节点和关系
    with open('data_rte_en.csv', 'r',encoding='utf-8') as f:
        reader = csv.reader(f)
        data = list(reader)
    print(data[1])  # 测试
    for i in range(1, len(data)):
        # node_s = (
        head_type = data[i][0]
        head = data[i][1]
        tail_type = data[i][3]
        tail = data[i][4]
        find_or_create_node(head_type, head)
        # node_s = Node(data[i][0], name=data[i][1])
        # node_o = Node(data[i][2], name=data[i][3])
        # node_o =
        find_or_create_node(tail_type, tail)
        # relation = Node('address', name=data[i][4])
        # relation1 = Node('sex', name=data[i][3])

        # graph.create(node_s)
        # graph.create(node_o)
        # graph.create(relation)
        # graph.create(relation1)
        relation_type = data[i][2]
        s = f"MATCH (head: {head_type}),(tail:{tail_type}) WHERE head.name = '" + head + "' AND tail.name = '" + tail + \
            "' CREATE (head)-[r:" + relation_type + "] -> (tail) RETURN r"
        # print(s)
        try:
            graph.run(s)
        except:
            pass
        graph.run(s)
        # relate = Relationship(node_s, data[i][2], node_o)
        # sex = Relationship(node, '性别', relation1)
        # graph.create(relate)
        # graph.create(sex)

    print("创建节点和关系成功")
    # endregion

    # region 导出所有实体
    label_name = ['Check', 'Department', 'Disease', 'Drug', 'Food', 'Symptom']
    entity_name_list = [[] for i in range(len(label_name))]
    # 执行查询
    # 定义 Cypher 查询语句
    query = """
        MATCH (n)
        RETURN DISTINCT labels(n) AS labels, n.name AS entity
        """
    results = graph.run(query)
    for record in results:
        labels = record["labels"][0]
        entity_name = record["entity"]
        for j in range(len(label_name)):
            if labels == label_name[j]:
                entity_name_list[j].append(entity_name)
                break

    for i in range(len(label_name)):
        label = label_name[i]
        file_name = f"./data/entities/{label}.txt"
        with open(file_name, "w", encoding="utf-8") as file:
            pass
        with open(file_name, "a", encoding="utf-8") as file:
            for item in entity_name_list[i]:
                file.write(f"{item}\n")  # 写入词库

if __name__ == "__main__":
    main()
